import numpy as np
import torch
import math
import os
import logging

import matplotlib
from mpmath.libmp import fhalf

matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib import animation 

import random

'''
Some auxiliary functions to be used in the experiment, such as reward adjustment, reading trajectory and so on
'''

def str2bool(v):
    if isinstance(v, bool):
        return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        False


class Log:
    def __init__(self, file_path=None, level=logging.INFO):
        # logging.basicConfig(
        #     level=logging.DEBUG,
        #     # format='%(asctime)s - %(levelname)s - %(message)s',
        #     format='%(message)s',
        #     filename=file_path,
        #     filemode='a',
        # )

        logger = logging.getLogger()
        logger.setLevel(level)

        if file_path is None:
            # 输出到控制台
            log_h = logging.StreamHandler()
        else:
            # 输出到文件
            log_h = logging.FileHandler(file_path, encoding='utf-8')

        log_h.setLevel(level)
        formatter = logging.Formatter('%(message)s')
        log_h.setFormatter(formatter)

        logger.addHandler(log_h)

        self.logger = logger


    def close(self):
        pass

    def record(self, s):
        # logging.info(s)
        print(s)
        self.logger.info(s)

    def debug(self, s):
        # logging.debug(s)
        self.logger.debug(s)


# Smooth drawing image
def plotSmoothAndSaveFig(interval, data, filePath):
    temp_list = []
    if len(data) > interval:
        for index in range(interval, len(data)):
            temp_list.append(np.average(np.array(data[index - interval:index])))
    else:
        for index in range(min(interval, len(data))):
            temp_list.append(data[index])
    # plt.figure()
    plt.plot(temp_list)
    plt.savefig(filePath)
    plt.close()


def display_frames_as_gif(frames, id, path):
    try:
        dpi = 60
        h, w = frames[0].shape[:2]
        figsize = (w / dpi, h / dpi)
        plt.figure(figsize=figsize, dpi=dpi)
        plt.axis('off')
        patch = plt.imshow(frames[0])

        def animate(i):
            patch.set_data(frames[i])

        anim = animation.FuncAnimation(plt.gcf(), animate, frames=len(frames), interval=5)
        plt.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=None, hspace=None)
        anim.save(os.path.join(path, f'{id}.gif'), writer='pillow', fps=30)
        plt.close()
    except Exception as e:
        print(f'env::display_frames_as_gif error -> {e}')

